"""仅适用于简单圆柱投影的图像数据，将源图像拆成训练集所用子块，同时又将子块所检测结果反推回原图像位置产生经纬度标注文件"""

import numpy as np
import tqdm
import os
import argparse
import yaml
import importlib
from itertools import product
import cv2
import copy


def arg_parser():
    arg_parse = argparse.ArgumentParser()
    arg_parse.add_argument(
        "-c",
        "--config_path",
        type=str,
        default="/home/a804_cbf/Code/luner_crater/config/2080x4/HG-HV.yaml",
        help="The path of the configuration file in yaml format",
    )
    arg_parse.add_argument(
        "--read",
        action="store_true",
        help="Read DEM and NAC image and save subfigs",
    )
    arg_parse.add_argument(
        "--write",
        action="store_true",
        help="write the labeled crater to a csv file",
    )
    arg_parse.add_argument(
        "--test",
        action="store_true",
        help="Test the labeled crater on the image",
    )
    return arg_parse.parse_args()


class Projection:
    def __init__(self, *args, image=None, **kwargs) -> None:
        self._image = image

    def forward(self, lat, lon):
        raise NotImplementedError

    def inverse(self, x, y):
        raise NotImplementedError

    def update(self, **kwargs):
        for key, value in kwargs.items():
            if hasattr(self, "_" + key):
                setattr(self, "_" + key, value)
        return self

    def __setitem__(self, key, value):
        self._image[key] = value

    def __getitem__(self, index: slice):
        """
        支持裁剪后，投影坐标系也相应改变，仅需要改变输入的序列值，程序将自动返回更改坐标系原点后的投影坐标
        """
        assert self._image is not None, "The image is not provided"
        assert len(index) == 2, "The index should be a 2D slice"
        lat0, lon0 = self.inverse(x=index[0].start, y=index[1].start)
        return copy.deepcopy(self).update(
            lat0=lat0, lon0=lon0, image=self._image[index]
        )

    @property
    def image(self):
        return self._image

    def __repr__(self) -> str:
        return f"Projection({self._image})"


class Equirectangular(Projection):
    def __init__(
        self,
        lat0,
        lon0,
        lat1,
        radius,
        pixel_size,
        image=None,
        **kwargs,
    ):
        """
        等距圆柱投影方式，虽然名义上投影按中心经线和标准纬线进行，实际上对应于图上时，中心经线等于起始经线，标准纬线等于起始纬线，平移量offset为0，这样才能保证投影前后，图像起始点与投影起始点一致 \\
        Arguments:
            lat1 (float) : 标准纬线，单位为度
            lon0 (float) : 中心经线，单位为度
            radius (float) : 地球半径，单位为米
            pixel_size (float) : 像素大小，单位为米/像素
        """
        super().__init__(image=image)
        self._lat1 = lat1
        self._lat0 = lat0
        self._lon0 = lon0
        self._pixel_size_radius_ratio = pixel_size / radius * 180 / np.pi

    def forward(self, lat, lon):
        """
        Arguments:
            lat (float) : 纬度，单位为度
            lon (float) : 经度，单位为度
        """
        x = (self._lat0 - lat) / self._pixel_size_radius_ratio
        y = (
            (lon - self._lon0)
            * np.cos(self._lat1 * np.pi / 180)
            / self._pixel_size_radius_ratio
        )
        return x, y

    def inverse(self, x, y):
        """
        Arguments:
            x (float) : 纵向坐标，单位为像素
            y (float) : 横向坐标，单位为像素
        """
        lat = self._lat0 - x * self._pixel_size_radius_ratio
        lon = (
            y / np.cos(self._lat1 * np.pi / 180) * self._pixel_size_radius_ratio
            + self._lon0
        )
        return lat, lon


class SimpleCylinder(Projection):
    def __init__(
        self, lat0, lon0, lat_range, lon_range, x_size, y_size, image=None, **kwargs
    ):
        """
        说明：
        Arguments:
            lat0 (float) : 中心纬度
            lon0 (float) : 中心经度
            lat_range (float) : 纬度范围
            lon_range (float) : 经度范围
            x_size (int) : 图像高度
            y_size (int) : 图像宽度
        """
        super().__init__(image=image)
        self._lat0 = lat0
        self._lon0 = lon0
        self._pixel_x_size = x_size / lat_range
        self._pixel_y_size = y_size / lon_range

    def forward(self, lat, lon):
        x = (self._lat0 - lat) * self._pixel_x_size
        y = (lon - self._lon0) * self._pixel_y_size
        return x, y

    def inverse(self, x, y):
        lat = self._lat0 - x / self._pixel_x_size
        lon = y / self._pixel_y_size + self._lon0
        return lat, lon


def from_IMG(config, base_batch, lats_clip, lons_clip, **kwargs):
    START_BYTE = config["start"]
    LINES = config["lines"]
    LINE_SAMPLES = config["samples"]
    BYTES_PER_PIXEL = config["byte"]
    module, cls = config["dtype"].split(".")
    dtype = getattr(importlib.import_module(module, package=None), cls)
    image_path = config["path"]
    radius = config["radius"]
    lon0 = config["lons"][0]
    lat0 = config["lats"][0]
    lat1 = config["lat1"]
    pixel_size = config["pixel_size"]
    num_batches = LINES // base_batch
    erp = Equirectangular(lat0, lon0, lat1, radius, pixel_size)

    x1, y1 = erp.forward(lats_clip[0], lons_clip[0])
    x2, y2 = erp.forward(lats_clip[1], lons_clip[1])
    x1, y1, x2, y2 = round(x1), round(y1), round(x2), round(y2)
    interest_area = np.zeros((x2 - x1, y2 - y1), dtype=np.float32)
    # 获取最大值与最小值
    max_value = -np.inf
    min_value = np.inf
    with open(image_path, "rb") as f:
        # 跳过前LABEL_BYTES的标签字节
        f.read(START_BYTE)
        for row in tqdm.tqdm(range(num_batches), desc="Preparing data"):
            # 读取一个标准batch的数据，即960行，184320列，共计192个batch
            data = f.read(LINE_SAMPLES * BYTES_PER_PIXEL * base_batch)
            # 将数据转换为numpy数组
            data_array = np.frombuffer(data, dtype=dtype).reshape(
                base_batch, LINE_SAMPLES
            )
            img_array = data_array.copy()
            img_array[img_array == img_array.min()] = 0
            if img_array.max() > max_value:
                max_value = img_array.max()
            if img_array.min() < min_value:
                min_value = img_array.min()
    with open(image_path, "rb") as f:
        # 跳过前LABEL_BYTES的标签字节
        f.read(START_BYTE)
        for row in tqdm.tqdm(range(num_batches), "Reading data"):
            # 读取一个标准batch的数据，即960行，184320列，共计192个batch
            data = f.read(LINE_SAMPLES * BYTES_PER_PIXEL * base_batch)
            # 将数据转换为numpy数组
            data_array = np.frombuffer(data, dtype=dtype).reshape(
                base_batch, LINE_SAMPLES
            )

            up = row * base_batch
            down = up + base_batch
            # 退出
            if up > x2 or down < x1:
                continue
            # 求出(left, up)和(right, down)对应的矩形与(x1, y1)和(x2, y2)对应的矩形的交集
            up = max(up, x1)
            down = min(down, x2)
            right = down % base_batch
            if right == 0:
                right = base_batch
            interest_area[up - x1 : down - x1] = data_array[
                up % base_batch : right, y1:y2
            ]

        img_array = interest_area.copy()
        img_array[img_array == img_array.min()] = min_value
        img_array = (img_array - min_value) / (max_value - min_value) * 255
    return erp.update(lat0=lats_clip[0], lon0=lons_clip[0], image=img_array)


def from_TIFF(path, lats, lons, lines, samples, lats_clip, lons_clip, **kwargs):
    LATITUDE_RANGE = lats[0] - lats[1]
    LONGITUDE_RANGE = lons[1] - lons[0]

    scp = SimpleCylinder(
        lats[0], lons[0], LATITUDE_RANGE, LONGITUDE_RANGE, lines, samples
    )
    img = cv2.imread(path, cv2.IMREAD_UNCHANGED)
    x1, y1 = scp.forward(lats_clip[0], lons_clip[0])
    x2, y2 = scp.forward(lats_clip[1], lons_clip[1])
    x1, y1, x2, y2 = round(x1), round(y1), round(x2), round(y2)
    return scp.update(lat0=lats_clip[0], lon0=lons_clip[0], image=img[x1:x2, y1:y2])


def to_windows(
    nac_prj: Projection,
    root_dir,
    base_batch,
    window,
    is_labeled=False,
    label_path=None,
    **kwargs,
):
    if not os.path.exists(os.path.join(root_dir, "nac")):
        os.makedirs(os.path.join(root_dir, "nac"))
    if is_labeled:
        assert label_path is not None, "The label path is not provided"
        labels = []
        with open(label_path, "r") as f:
            f.readline()
            for line in f:
                lon, lat, d = map(float, line.strip().split(",")[1:4])
                r = d * 1000 / 2.4
                x, y = nac_prj.forward(lat, 360 + lon)
                if 0 <= x < nac_prj.image.shape[0] and 0 <= y < nac_prj.image.shape[1]:
                    labels.append((y - r, x - r, y + r, x + r))
    # 应当考虑dem和nac分辨率不相同的情况，此时需要以经纬度的方式进行对齐
    # 通常DEM分辨率小于NAC，这里的window和base_batch应当以DEM为准
    # 以滑动窗为50x50，步长为50，保存为子块
    for i, j in tqdm.tqdm(
        product(
            range(0, nac_prj.image.shape[1] - base_batch, window),
            range(0, nac_prj.image.shape[0] - base_batch, window),
        ),
        "Saving subfigs",
    ):
        batch = nac_prj.image[j : j + base_batch, i : i + base_batch]
        if batch.shape[0] != base_batch or batch.shape[1] != base_batch:
            continue
        if is_labeled:
            with open(
                os.path.join(root_dir, "nac", f"{i}_{j}_{base_batch}.txt"), "w"
            ) as f:
                flag = False
                for x1, y1, x2, y2 in labels:
                    if (
                        x1 > i
                        and y1 > j
                        and x2 < i + base_batch
                        and y2 < j + base_batch
                    ):
                        f.write(f"{x1 - i},{y1 - j},{x2 - i},{y2 - j}\n")
                        flag = True
                f.flush()
            if flag:
                cv2.imwrite(
                    os.path.join(root_dir, "nac", f"{i}_{j}_{base_batch}.png"),
                    batch,
                )
            else:
                os.remove(os.path.join(root_dir, "nac", f"{i}_{j}_{base_batch}.txt"))
        else:
            cv2.imwrite(
                os.path.join(root_dir, "nac", f"{i}_{j}_{base_batch}.png"),
                batch,
            )


def get_crater(x1, y1, x2, y2, prj: Projection, loose=0.2):
    # 加入松弛因子，以确保获取完整的陨石坑
    x1 = round(max(x1 - loose * (x2 - x1), 0))
    x2 = round(min(x2 + loose * (x2 - x1), prj.image.shape[0]))
    y1 = round(max(y1 - loose * (y2 - y1), 0))
    y2 = round(min(y2 + loose * (y2 - y1), prj.image.shape[1]))
    return prj[x1:x2, y1:y2]


def to_latlons(nac_prj: Projection, dem_prj: Projection, points_path):
    with open(points_path, "r") as f:
        file = os.path.basename(points_path).split(".")[0]
        row, col, base_batch = list(map(int, file.split(".")[0].split("_")))
        # 舍去第一行
        f.readline()
        # 读取检测框，格式为xyxy，是NAC的像素坐标
        for line in f:
            x_nac = list(map(float, line.strip().split(",")))
            # score = x_nac[-1]
            # if score < 0.5:
            #     continue
            x_nac = [x_nac[0] + row, x_nac[1] + col, x_nac[2] + row, x_nac[3] + col]
            # 排除小陨石坑
            if x_nac[2] - x_nac[0] < 10 or x_nac[3] - x_nac[1] < 10:
                continue
            cnty = (x_nac[0] + x_nac[2]) / 2
            cntx = (x_nac[1] + x_nac[3]) / 2
            r = (x_nac[2] - x_nac[0] + x_nac[3] - x_nac[1]) / 4
            lat, lon = nac_prj.inverse(cntx, cnty)
            yield lat, lon, r


def to_labels(nac_prj: Projection, dem_prj: Projection, points_path):
    """
    根据不同的子块序号，将其转为经纬度坐标，输入的应该是子块内检测到的每个陨石坑中心坐标、子块的横纵值
    """
    with open(points_path, "r") as f:
        file = os.path.basename(points_path).split(".")[0]
        row, col, base_batch = list(map(int, file.split(".")[0].split("_")))
        # 舍去第一行
        f.readline()
        # 读取检测框，格式为xyxy，是NAC的像素坐标
        for line in f:
            x_nac = list(map(float, line.strip().split(",")))
            # score = x_nac[-1]
            # if score < 0.5:
            #     continue
            x_nac = [x_nac[1] + col, x_nac[0] + row, x_nac[3] + col, x_nac[2] + row]
            # 排除小陨石坑
            if x_nac[2] - x_nac[0] < 10 or x_nac[3] - x_nac[1] < 10:
                continue
            lat1, lon1 = nac_prj.inverse(x_nac[0], x_nac[1])
            lat2, lon2 = nac_prj.inverse(x_nac[2], x_nac[3])
            # 求出DEM图像下的坐标
            x_dem1, y_dem1 = dem_prj.forward(lat1, lon1)
            x_dem2, y_dem2 = dem_prj.forward(lat2, lon2)
            # x1,y1,x2,y2 均为nac图像内的坐标，此时应当求出对应位置处的DEM图像坐标
            crater_dem = get_crater(x_dem1, y_dem1, x_dem2, y_dem2, dem_prj)
            x_nac = [x_nac[0], x_nac[1] + 15, x_nac[2], x_nac[3] + 15]
            # 获取陨石坑图像
            crater_nac = get_crater(*x_nac, nac_prj)
            yield crater_nac, crater_dem


def to_global(dem_prj: Projection, nac_prj: Projection):
    dem_image = dem_prj.image
    nac_image = nac_prj.image
    dem_image[dem_image == dem_image.min()] = dem_image.max()
    dem_image = (
        (dem_image - dem_image.min()) / (dem_image.max() - dem_image.min()) * 255
    )
    nac_image = (
        (nac_image - nac_image.min()) / (nac_image.max() - nac_image.min()) * 255
    )
    dem_image = cv2.cvtColor(dem_image, cv2.COLOR_GRAY2BGR)
    nac_image = cv2.cvtColor(nac_image, cv2.COLOR_GRAY2BGR)
    with open(
        "/disk527/DataDisk/a804_cbf/datasets/lunar_crater/chang_e/Liu_chang_e_5_annotation.txt",
        "r",
    ) as f:
        f.readline()
        with open("craters.csv", "w") as w:
            for line in f:
                lon, lat, d = list(map(float, line.strip().split(",")[1:4]))
                x, y = dem_prj.forward(lat, 360 + lon)
                cv2.circle(
                    dem_image, (round(y), round(x)), round(d * 1000 / 3), (0, 0, 255), 1
                )
                x, y = nac_prj.forward(lat, 360 + lon)
                cv2.circle(
                    nac_image,
                    (round(y), round(x)),
                    round(d * 1000 / 2.4),
                    (0, 0, 255),
                    1,
                )
                if 0 <= x <= nac_image.shape[0] and 0 <= y <= nac_image.shape[1]:
                    w.write(f"{lat},{360+lon},{d/2*1000}\n")
    cv2.imwrite("dem-craters.png", dem_image)
    cv2.imwrite("nac-craters.png", nac_image)


if __name__ == "__main__":
    args = arg_parser()
    with open(args.config_path, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
        nac_config = config["nac"]
        dem_config = config["dem"]
        subfig_config = config["subfig"]
    dem_prj = from_TIFF(**dem_config, **subfig_config)
    nac_prj = from_IMG(nac_config, **subfig_config)
    dem_image = dem_prj.image
    nac_image = nac_prj.image
    output_dir = subfig_config["root_dir"]
    if args.read:
        to_windows(nac_prj, **subfig_config)
    if args.write:
        # dem_image[dem_image == dem_image.min()] = dem_image.max()
        # dem_image = (
        #     (dem_image - dem_image.min()) / (dem_image.max() - dem_image.min()) * 255
        # )
        nac_image = (
            (nac_image - nac_image.min()) / (nac_image.max() - nac_image.min()) * 255
        )
        # dem_image = cv2.cvtColor(dem_image, cv2.COLOR_GRAY2BGR)
        nac_image = cv2.cvtColor(nac_image, cv2.COLOR_GRAY2BGR)
        dem_prj.update(image=dem_image)
        nac_prj.update(image=nac_image)
        with open(os.path.join(output_dir, "crater_catalog.txt"), "w") as f:
            files = os.listdir(
                os.path.join("/disk527/sdb1/a804_cbf/datasets/chang_e/bboxes")
            )
            for file in map(lambda x: x.split(".")[0], files):
                points_path = os.path.join(
                    "/disk527/sdb1/a804_cbf/datasets/chang_e/bboxes", file + ".txt"
                )
                nac_prj, dem_prj = to_labels(nac_prj, dem_prj, points_path)
        cv2.imwrite(os.path.join(output_dir, "dem.png"), dem_prj.image)
        cv2.imwrite(os.path.join(output_dir, "nac.png"), nac_prj.image)
    if args.test:
        to_global(dem_prj, nac_prj)
